#include "kernel_operator.h"
#include "op_common.h"
#include <cstdint>

using namespace AscendC;

template <typename _DT_X, typename _DT_Y> class Kernel {
public:
  TPipe pipe;
  DefInTensor(X);
  DefOutTensor(Y);

  DefBufVECIN(WORK_LOCAL);
  DefBufVECIN(CASTED_X);
  DefOutQue(SUM_VALUE);

  int64_t nxc;
  int64_t d;
  int64_t tileLength;
  int64_t finalLength;
  TypeOf(X) c = 0;
  int64_t current_j = 0;

public:
  __aicore__ inline Kernel() {}

  template <typename T>
  __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, T &tiling_data) {
    this->nxc = tiling_data.nxc;
    this->d = tiling_data.d;
    this->tileLength = tiling_data.tileLength;
    this->finalLength = tiling_data.finalLength;
    // 切分global
    GSetBuffer(X, x, 0, nxc * d);
    GSetBuffer(Y, y, 0, nxc);
    // 初始化队列
    InitQueueSimple(X, tileLength);
    InitQueue(Y, 32);
    InitQueue(SUM_VALUE, 32);
    // 初始化BUF
    auto work_locallenth = workLocalSize<TypeOf(X)>(tileLength);
    InitTBufBuffer(WORK_LOCAL, work_locallenth * sizeof(TypeOf(X)));
    // InitTBufBuffer(WORK_LOCAL, work_locallenth * sizeof(float));
    // if constexpr (std::is_same_v<TypeOf(X), half>)
    // {
    //     InitTBufBuffer(CASTED_X, tileLength * sizeof(float));
    // }
  }

  __aicore__ inline void Process() {
    auto loopCount = CEIL_DIV(d, tileLength);
    auto finnal_progress = loopCount - 1;
    for (auto i = 0; i < nxc; ++i) {
      GTensorName(Y).SetValue(i, TypeOf(Y)(0));
    }
    for (auto j = 0; j < nxc; ++j) {
      for (uint32_t i = 0; i < finnal_progress; ++i) {
        CopyIn(i, j, tileLength);
        Compute(tileLength, j);
        CopyOut(i, j, tileLength);
      }
      CopyInFinal(finnal_progress, j, finalLength);
      Compute(finalLength, j);
      CopyOutFinal(finnal_progress, j, finalLength);
    }
  }

  __aicore__ inline void CopyIn(uint32_t i, uint32_t j, uint32_t calcCount) {
    EnQueGlobal2Local(X, j * d + i * tileLength, tileLength);
    EnQueGlobal2Local(Y, j, elementsPerBlock<TypeOf(X)>());
  }

  __aicore__ inline void CopyInFinal(uint32_t i, uint32_t j,
                                     uint32_t calcCount) {
    EnQueGlobal2Local(X, j * d + i * tileLength,
                      ALIGN_TO(finalLength, elementsPerBlock<TypeOf(X)>()));
    EnQueGlobal2Local(Y, j, elementsPerBlock<TypeOf(X)>());
  }

  __aicore__ inline void Compute(uint32_t calcCount, int64_t j) {
    if constexpr (std::is_same_v<TypeOf(X), float>) {
      DeQueSimple(X);
      DeQueSimple(Y);
      TBufGet(WORK_LOCAL, TypeOf(X));
      QueAlloc(SUM_VALUE, TypeOf(X));
      ReduceSum(LTensor(SUM_VALUE), LTensor(X), BTensor(WORK_LOCAL), calcCount);
      Add(LTensor(SUM_VALUE), LTensor(SUM_VALUE), LTensor(Y), 1);
      EnQue(SUM_VALUE);
      QueFree(X);
      QueFree(Y);
    } else if constexpr (std::is_same_v<TypeOf(X), half>) {
      if (current_j != j) {
        c = 0;
      }

      DeQueSimple(Y);
      DeQueSimple(X);
      TBufGet(WORK_LOCAL, TypeOf(X));
      QueAlloc(SUM_VALUE, TypeOf(X));
      TypeOf(X) sumv = LTensor(Y)(0);
      int16_t v_int16 = *((int16_t *)&sumv);
      v_int16 = v_int16 & 0x7fff;
      if (v_int16 != infInt16) {
        print("v_int16 = %d\n", v_int16);
        for (auto i = 0; i < calcCount; ++i) {
          float y = float(LTensor(X)(i)) - float(c);
          float t = float(sumv) + y;
          c = (t - float(sumv)) - y;
          if (t < -65504.0f) {
            sumv = ninfHalf;
            break;
          } else if (t > 65504.0f) {
            sumv = infHalf;
            break;
          } else {
            sumv = t;
          }
        }
      }
      print("sumv = %f\n", sumv);
      LTensor(SUM_VALUE)(0) = sumv;
      EnQue(SUM_VALUE);
      QueFree(X);
      QueFree(Y);
    }
  }

  __aicore__ inline void CopyOut(uint32_t i, uint32_t j, uint32_t calcCount) {
    DeQue(SUM_VALUE, TypeOf(X));
    print("CopyOut = %f\n", LTensor(SUM_VALUE)(0));
    GTensor(Y)(j) = LTensor(SUM_VALUE)(0);
    QueFree(SUM_VALUE);
  }

  __aicore__ inline void CopyOutFinal(uint32_t i, uint32_t j,
                                      uint32_t calcCount) {
    if constexpr (std::is_same_v<TypeOf(X), float>) {
      DeQue(SUM_VALUE, TypeOf(X));
      GTensor(Y)(j) = LTensor(SUM_VALUE)(0) / float(d);
      QueFree(SUM_VALUE);
    } else {
      DeQue(SUM_VALUE, TypeOf(X));
      auto v = LTensor(SUM_VALUE)(0);
      int16_t v_int16 = *((int16_t *)&v);
      if (v_int16 == infInt16) {
        GTensor(Y)(j) = infHalf;
      } else if (v_int16 == ninfInt16) {
        GTensor(Y)(j) = ninfHalf;
      } else {
        GTensor(Y)(j) = float(v) / float(d);
      }
      QueFree(SUM_VALUE);
    }
  }
};

extern "C" __global__ __aicore__ void
global_avg_pool(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
  GET_TILING_DATA(tiling_data, tiling);
  Kernel<DTYPE_X, DTYPE_Y> op;
  op.Init(x, y, tiling_data);
  op.Process();
}